mlx 0.32.1 - #299369

## What's Changed

Fix int64 type cast error when loading GGUF metadata arrays by @danlee2002 in Fix int64 type cast error when loading GGUF metadata arrays ml-explore/mlx#3823
Document default values in normalization layer docstrings by @Pablosinyores in Document default values in normalization layer docstrings ml-explore/mlx#3819
Warn at configure time when NAX kernels are disabled by @pierre427 in Warn at configure time when NAX kernels are disabled ml-explore/mlx#3824
[CUDA][Improvement] RMSNorm forward speed up by @nastya236 in [CUDA][Improvement] RMSNorm forward speed up ml-explore/mlx#3850
Fix captured random state in compile by @angeloskath in Fix captured random state in compile ml-explore/mlx#3828
Fix JIT preamble header filter matching project paths containing "Xcode" by @apocryphx in Fix JIT preamble header filter matching project paths containing "Xcode" ml-explore/mlx#3873
Document default value of p in dropout layer docstrings by @ayaangazali in Document default value of p in dropout layer docstrings ml-explore/mlx#3870
[WIP] [CUDA] fsdp by @nastya236 in [WIP] [CUDA] fsdp ml-explore/mlx#3768
Fix triplet_loss docstring to document the reduced output shape by @ayaangazali in Fix triplet_loss docstring to document the reduced output shape ml-explore/mlx#3884
Zero-copy CPU import: mx.array(host_buffer, copy=False) on unified memory by @HaoXuAI in Zero-copy CPU import: mx.array(host_buffer, copy=False) on unified memory ml-explore/mlx#3872
Round MLX_SDPA_BLOCKS up to a multiple of 32 by @pierre427 in Round MLX_SDPA_BLOCKS up to a multiple of 32 ml-explore/mlx#3875
Reuse Metal WAR tracking hash tables by @neilmehta24 in Reuse Metal WAR tracking hash tables ml-explore/mlx#3882
Use unroll_count(4) for the NAX attention Q@K.T loop by @wyanzhao in Use unroll_count(4) for the NAX attention Q@K.T loop ml-explore/mlx#3843
metal: add gemv_wide for fp16/bf16 matmuls of a few rows by @jessegross in metal: add gemv_wide for fp16/bf16 matmuls of a few rows ml-explore/mlx#3888
Fix broken docstring rendering in Linear and RNN by @ayaangazali in Fix broken docstring rendering in Linear and RNN ml-explore/mlx#3890
Fix Adamax betas docstring and MultiOptimizer filters type by @ayaangazali in Fix Adamax betas docstring and MultiOptimizer filters type ml-explore/mlx#3889
Fix incorrect nvfp4 quantized_matmul through the split-K path by @metascroy in Fix incorrect nvfp4 quantized_matmul through the split-K path ml-explore/mlx#3854
[Metal] Avoid regex in custom kernel name generation by @aleroot in [Metal] Avoid regex in custom kernel name generation ml-explore/mlx#3869
Fix prod dtype promotion when reducing a size-1 axis by @eyupcanakman in Fix prod dtype promotion when reducing a size-1 axis ml-explore/mlx#3898
[CUDA] Fix grid overflow in gemm conv unfold kernels for >= 65,536 output positions by @AdamDLuz in [CUDA] Fix grid overflow in gemm conv unfold kernels for >= 65,536 output positions ml-explore/mlx#3893
[CUDA] columnwise quantize with tma by @nastya236 in [CUDA] columnwise quantize with tma ml-explore/mlx#3157
metal: reduce NVFP4 scales per 16-lane group by @jessegross in metal: reduce NVFP4 scales per 16-lane group ml-explore/mlx#3934
Update homebrew in CI by @angeloskath in Update homebrew in CI ml-explore/mlx#3946
Make index autodiff errors explicitly recommend stop_gradient by @dogukanveziroglu in Make index autodiff errors explicitly recommend stop_gradient ml-explore/mlx#3820
Making JACCL coordinator optional by @angeloskath in Making JACCL coordinator optional ml-explore/mlx#3899
Fix docstring mismatches in the Python bindings by @ayaangazali in Fix docstring mismatches in the Python bindings ml-explore/mlx#3948
Raise a clear error for an invalid quantization mode in nn layers by @ayaangazali in Raise a clear error for an invalid quantization mode in nn layers ml-explore/mlx#3914
Fix incorrect examples and outputs in the usage docs by @ayaangazali in Fix incorrect examples and outputs in the usage docs ml-explore/mlx#3956
Skip test_gather_qmm_sorted cpu test on M1 mac by @zcbenz in Skip test_gather_qmm_sorted cpu test on M1 mac ml-explore/mlx#3973
Fixes an axis mismatch bug in matrix norm for case -1 and 1  by @danlee2002 in Fixes an axis mismatch bug in matrix norm for case -1 and 1  ml-explore/mlx#3827
Fix BatchNorm running variance estimator by @ishtihoss in Fix BatchNorm running variance estimator ml-explore/mlx#3817
Fix build error caused by TMA macro guard by @zcbenz in Fix build error caused by TMA macro guard ml-explore/mlx#3988
Fix Glorot/He uniform init docstrings to label the uniform bound, not sigma by @vineethsaivs in Fix Glorot/He uniform init docstrings to label the uniform bound, not sigma ml-explore/mlx#3831
Fix custom metal kernel cache collision for same name, different source by @katlun-lgtm in Fix custom metal kernel cache collision for same name, different source ml-explore/mlx#3833
Fix log_cosh_loss docstring to document the element-wise loss by @winklemad in Fix log_cosh_loss docstring to document the element-wise loss ml-explore/mlx#3846
Fix step activation docstring to match >= threshold behavior by @ayaangazali in Fix step activation docstring to match >= threshold behavior ml-explore/mlx#3902
docs: document the reduced-precision float32 default and MLX_ENABLE_TF32 by @stoyoda0012-cyber in docs: document the reduced-precision float32 default and MLX_ENABLE_TF32 ml-explore/mlx#3894
Fix InstanceNorm Shape docstring to require at least 3 dimensions by @ayaangazali in Fix InstanceNorm Shape docstring to require at least 3 dimensions ml-explore/mlx#3903
Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn by @ayaangazali in Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn ml-explore/mlx#3906
Export C++20 requirement to CMake consumers by @PhysicistJohn in Export C++20 requirement to CMake consumers ml-explore/mlx#3971
Fix implicit thread address space qualifier becoming explicit in metal 4.1  by @louen in Fix implicit thread address space qualifier becoming explicit in metal 4.1  ml-explore/mlx#3963
Fix Transformer ignoring a custom encoder or decoder with no parameters by @ayaangazali in Fix Transformer ignoring a custom encoder or decoder with no parameters ml-explore/mlx#3962
docs: remove references to removed --no-verify-script launch flag by @latent-9 in docs: remove references to removed --no-verify-script launch flag ml-explore/mlx#3959
added eye(0) support and tests by @aaishwarymishra in added eye(0) support and tests ml-explore/mlx#3952
Refactor the JACCL ring and add threads for the multiple rings by @angeloskath in Refactor the JACCL ring and add threads for the multiple rings ml-explore/mlx#3900
Fix shapeless matmul with dynamic batch dimensions by @varshneydevansh in Fix shapeless matmul with dynamic batch dimensions ml-explore/mlx#3813
Fix JIT build with old macOS SDK by @metascroy in Fix JIT build with old macOS SDK ml-explore/mlx#3853
Fix cholesky_inv documented argument name to match the binding by @ayaangazali in Fix cholesky_inv documented argument name to match the binding ml-explore/mlx#3950
Add new_thread_unsafe_stream to the devices and streams docs by @ayaangazali in Add new_thread_unsafe_stream to the devices and streams docs ml-explore/mlx#3968
Add gather_qqmm by @zcbenz in Add gather_qqmm ml-explore/mlx#3757
Fix missing printoptions doc page from autosummary filename collision by @ayaangazali in Fix missing printoptions doc page from autosummary filename collision ml-explore/mlx#3985
Document ThreadLocalStream and iinfo in the API reference by @ayaangazali in Document ThreadLocalStream and iinfo in the API reference ml-explore/mlx#3986
Make "stop" optional in arange by @aaishwarymishra in Make "stop" optional in arange ml-explore/mlx#3982
Add softsign to the nn functions docs by @ayaangazali in Add softsign to the nn functions docs ml-explore/mlx#3989
Fix broken all_sum and Group references in the data parallelism example by @ayaangazali in Fix broken all_sum and Group references in the data parallelism example ml-explore/mlx#3996
Fix ast.metal_kernel typo in the custom Metal kernels guide by @ayaangazali in Fix ast.metal_kernel typo in the custom Metal kernels guide ml-explore/mlx#3997
Fix unresolved mx.array docstring references by @ayaangazali in Fix unresolved mx.array docstring references ml-explore/mlx#3990
python: fix bfloat16 buffer format itemsize mismatch by @reckylurker in python: fix bfloat16 buffer format itemsize mismatch ml-explore/mlx#3975
Template Metal complex scalar lanes by @PhysicistJohn in Template Metal complex scalar lanes ml-explore/mlx#3970
Pad 2D conv input channels to reach the specialized Metal kernel by @eyupcanakman in Pad 2D conv input channels to reach the specialized Metal kernel ml-explore/mlx#3904
Support head dimension 96 in Metal full attention by @dhiltgen in Support head dimension 96 in Metal full attention ml-explore/mlx#3943
Fix mx.remainder floored-mod for float16/bfloat16 on CPU by @sashko-zakharchuk in Fix mx.remainder floored-mod for float16/bfloat16 on CPU ml-explore/mlx#3976
Fix sorted gather_mm activation row stride by @metascroy in Fix sorted gather_mm activation row stride ml-explore/mlx#3960
Fix state corruption when a primitive throws during eval by @WindChimeRan in Fix state corruption when a primitive throws during eval ml-explore/mlx#3675
added complex support by @aaishwarymishra in added __complex__ support ml-explore/mlx#3984
Validate freeze and unfreeze keys against the whole model when recursing by @ayaangazali in Validate freeze and unfreeze keys against the whole model when recursing ml-explore/mlx#3966
Fix mlx.launch --python: flag is parsed but never forwarded to the launch script by @jonathan308 in Fix mlx.launch --python: flag is parsed but never forwarded to the launch script ml-explore/mlx#4002
Fix broken fully_shard reference in FullyShardedModule docstring by @ayaangazali in Fix broken fully_shard reference in FullyShardedModule docstring ml-explore/mlx#4007
Use the current interpreter in the comparative benchmark runner by @ayaangazali in Use the current interpreter in the comparative benchmark runner ml-explore/mlx#4016
Document MLX environment variables by @XXXXRT666 in Document MLX environment variables ml-explore/mlx#4000
Fix CUDA batched GEMV grid overflow by @jasp-nerd in Fix CUDA batched GEMV grid overflow ml-explore/mlx#3929
Fix all_gather benchmark collapsing its input to a scalar by @ayaangazali in Fix all_gather benchmark collapsing its input to a scalar ml-explore/mlx#4017
Check threadgroup size in the 1-pass sdpa_vector dispatch by @apocryphx in Check threadgroup size in the 1-pass sdpa_vector dispatch ml-explore/mlx#4018
Align example projects with the Python 3.10 minimum by @ayaangazali in Align example projects with the Python 3.10 minimum ml-explore/mlx#4024
Fix C++ benchmarks failing to build on overloaded astype by @ayaangazali in Fix C++ benchmarks failing to build on overloaded astype ml-explore/mlx#4025
Fix conv_transpose maxBufferLength failures on Metal via tiled unfold by @eyupcanakman in Fix conv_transpose maxBufferLength failures on Metal via tiled unfold ml-explore/mlx#3845
Give each host a unique rank in Hostfile.from_list by @ayaangazali in Give each host a unique rank in Hostfile.from_list ml-explore/mlx#4027
Fix crash when reporting partial rings in mlx.distributed_config by @ayaangazali in Fix crash when reporting partial rings in mlx.distributed_config ml-explore/mlx#4026
docs: Do not pass MLX_METAL_FAST_SYNCH=1 by default by @katlun-lgtm in docs: Do not pass MLX_METAL_FAST_SYNCH=1 by default ml-explore/mlx#4005
Fix signed-integer overflow in convolution shape arithmetic by @eyupcanakman in Fix signed-integer overflow in convolution shape arithmetic ml-explore/mlx#3938
Template Metal C2C FFT scalar lanes by @PhysicistJohn in Template Metal C2C FFT scalar lanes ml-explore/mlx#3969
Fix segfault in expand_dims for out of bounds negative axes by @Gusanidas in Fix segfault in expand_dims for out of bounds negative axes ml-explore/mlx#4021
Add optional dtype parameter to zeros_like and ones_like by @reckylurker in Add optional dtype parameter to zeros_like and ones_like ml-explore/mlx#4028
Fix mx.longsumexp output shape issue. by @JasonHonKL in Fix mx.longsumexp output shape issue. ml-explore/mlx#4030
Treat backend and envs as optional when parsing a hostfile by @ayaangazali in Treat backend and envs as optional when parsing a hostfile ml-explore/mlx#4039
Report bad hostfiles as CLI errors in mlx.launch by @ayaangazali in Report bad hostfiles as CLI errors in mlx.launch ml-explore/mlx#4040
Fix ceil error message to say ceil instead of floor by @PhysicistJohn in Fix ceil error message to say ceil instead of floor ml-explore/mlx#4042
Fix installed static MLX package on Windows by @yingjiacai in Fix installed static MLX package on Windows ml-explore/mlx#3848
Derive the qmv fast path K alignment from bits by @jlee5814 in Derive the qmv fast path K alignment from bits ml-explore/mlx#3965
Fix eigh UPLO and zero-size eigh/svd on the CPU by @kapellirohith in Fix eigh UPLO and zero-size eigh/svd on the CPU ml-explore/mlx#3834
Add metadata to exported functions by @eyupcanakman in Add metadata to exported functions ml-explore/mlx#3855
Fix empty matrix case in cholesky by @prady0t in Fix empty matrix case in cholesky ml-explore/mlx#4033
Reject a negative max_norm in clip_grad_norm by @ayaangazali in Reject a negative max_norm in clip_grad_norm ml-explore/mlx#4058
Build extensions with the interpreter running the build by @ayaangazali in Build extensions with the interpreter running the build ml-explore/mlx#4057
Reject negative dimensions in broadcast_to and random shapes by @ayaangazali in Reject negative dimensions in broadcast_to and random shapes ml-explore/mlx#4046
Report bad host arguments as CLI errors instead of tracebacks by @ayaangazali in Report bad host arguments as CLI errors instead of tracebacks ml-explore/mlx#4045
Fix Metal FFT for sizes above 2**20 by @adityasingh2400 in Fix Metal FFT for sizes above 2**20 ml-explore/mlx#4013
Request MSL 4.1 from the runtime compiler on macOS 27 by @magnuslundstedt in Request MSL 4.1 from the runtime compiler on macOS 27 ml-explore/mlx#4052
Fix concurrent Metal kernel cache lookup by @kitty-eu-org in Fix concurrent Metal kernel cache lookup ml-explore/mlx#4043
Refactor CPU ArgReduce dtype dispatch by @PhysicistJohn in Refactor CPU ArgReduce dtype dispatch ml-explore/mlx#4066
Refactor NumberOfElements dtype dispatch by @PhysicistJohn in Refactor NumberOfElements dtype dispatch ml-explore/mlx#4064
Refactor CPU Select dtype dispatch by @PhysicistJohn in Refactor CPU Select dtype dispatch ml-explore/mlx#4065
Raise qmv batch limit for large matrices on M5-class GPUs by @pierre427 in Raise qmv batch limit for large matrices on M5-class GPUs ml-explore/mlx#3791
Enable half-precision complex Metal FFT kernels by @PhysicistJohn in Enable half-precision complex Metal FFT kernels ml-explore/mlx#3981
Fix: mx.distributed.sum_scatter crashes on a scalar(0-d) by @JasonHonKL in Fix: mx.distributed.sum_scatter crashes on a scalar(0-d) ml-explore/mlx#4071
Return an empty result from pinv for zero-size inputs by @ayaangazali in Return an empty result from pinv for zero-size inputs ml-explore/mlx#4069
Handle zero-size inputs in eig and qr on the CPU by @ayaangazali in Handle zero-size inputs in eig and qr on the CPU ml-explore/mlx#4068
Normalize biases before encoding in gather_qmm_rhs by @dudududukim in Normalize biases before encoding in gather_qmm_rhs ml-explore/mlx#4056
Keep randint samples inside [low, high) by @adityasingh2400 in https://github.com/ml-explore/mlx/pull/4012
Allow saving empty arrays to npy and safetensors by @ayaangazali in Allow saving empty arrays to npy and safetensors ml-explore/mlx#4080
Fix DeviceType annotations for device and stream arguments by @XXXXRT666 in Fix DeviceType annotations for device and stream arguments ml-explore/mlx#4059
Refactor CPU DivMod dtype dispatch by @PhysicistJohn in Refactor CPU DivMod dtype dispatch ml-explore/mlx#4088
Refactor CPU SegmentedMM dtype dispatch by @PhysicistJohn in Refactor CPU SegmentedMM dtype dispatch ml-explore/mlx#4087
Refactor CPU Arange dtype dispatch by @PhysicistJohn in Refactor CPU Arange dtype dispatch ml-explore/mlx#4085
Fix Device and Stream lexicographic ordering by @wgu9 in Fix Device and Stream lexicographic ordering ml-explore/mlx#4086
Enable complex64 scatter addition on GPU by @PhysicistJohn in Enable complex64 scatter addition on GPU ml-explore/mlx#4078
Optimize large NVFP4 QMV on M5 Max by @dhiltgen in Optimize large NVFP4 QMV on M5 Max ml-explore/mlx#3961
ring: fail on peer disconnect instead of hanging forever by @erwinzhang7 in ring: fail on peer disconnect instead of hanging forever ml-explore/mlx#4060
Fix: propagate NaN when dealing with cummax and curmin by @JasonHonKL in Fix: propagate NaN when dealing with cummax and curmin ml-explore/mlx#4047
Refactor CPU equal_nan dtype dispatch by @PhysicistJohn in Refactor CPU equal_nan dtype dispatch ml-explore/mlx#4095
Refactor CPU Matmul dtype dispatch by @PhysicistJohn in Refactor CPU Matmul dtype dispatch ml-explore/mlx#4103
Refactor CPU Gather dtype dispatch by @PhysicistJohn in Refactor CPU Gather dtype dispatch ml-explore/mlx#4105
Refactor generic CPU unary dtype dispatch by @PhysicistJohn in Refactor generic CPU unary dtype dispatch ml-explore/mlx#4106
Refactor generic CPU binary dtype dispatch by @PhysicistJohn in Refactor generic CPU binary dtype dispatch ml-explore/mlx#4107
Validate pooling kernel size, stride and padding by @ayaangazali in Validate pooling kernel size, stride and padding ml-explore/mlx#4101
Only reject max and min when a reduced axis is empty by @ayaangazali in Only reject max and min when a reduced axis is empty ml-explore/mlx#4079
Fix SinusoidalPositionalEncoding silently overriding explicit scale=0.0 by @YassineMA03 in Fix SinusoidalPositionalEncoding silently overriding explicit scale=0.0 ml-explore/mlx#4098
Fix any and all treating -0.0 as nonzero by @ayaangazali in Fix any and all treating -0.0 as nonzero ml-explore/mlx#4090
Fix Eigh CPU dtype error message by @PhysicistJohn in Fix Eigh CPU dtype error message ml-explore/mlx#4130
Fix random distribution dtype error messages by @PhysicistJohn in Fix random distribution dtype error messages ml-explore/mlx#4131
Fix safetensors dtype error formatting by @PhysicistJohn in Fix safetensors dtype error formatting ml-explore/mlx#4132
Fix no-GPU synchronize error formatting by @PhysicistJohn in Fix no-GPU synchronize error formatting ml-explore/mlx#4133
Fix: mx.isinf doens't consider imaginary plane. by @JasonHonKL in Fix: mx.isinf doens't consider imaginary plane. ml-explore/mlx#4092
Refactor CPU Copy destination dtype dispatch by @PhysicistJohn in Refactor CPU Copy destination dtype dispatch ml-explore/mlx#4104
Refactor CPU GatherAxis dtype dispatch by @PhysicistJohn in Refactor CPU GatherAxis dtype dispatch ml-explore/mlx#4109
Refactor inexact CPU unary dtype dispatch by @PhysicistJohn in Refactor inexact CPU unary dtype dispatch ml-explore/mlx#4111
Refactor CPU comparison dtype dispatch by @PhysicistJohn in Refactor CPU comparison dtype dispatch ml-explore/mlx#4112
Replace Ring's local dtype macro with common dispatch by @PhysicistJohn in Replace Ring's local dtype macro with common dispatch ml-explore/mlx#4120
Refactor CPU ArgSort dtype dispatch by @PhysicistJohn in Refactor CPU ArgSort dtype dispatch ml-explore/mlx#4123
Refactor CPU Softmax dtype dispatch by @PhysicistJohn in Refactor CPU Softmax dtype dispatch ml-explore/mlx#4127
Use common dtype dispatch for Metal Arange by @PhysicistJohn in Use common dtype dispatch for Metal Arange ml-explore/mlx#4128
Pick BM from rows per expert in gather_qmm_rhs_nax by @dwijenpatel in Pick BM from rows per expert in gather_qmm_rhs_nax ml-explore/mlx#4023
Enable complex64 take_along_axis backward on Metal by @PhysicistJohn in Enable complex64 take_along_axis backward on Metal ml-explore/mlx#4094
Add mx.searchsorted with CPU, Metal and CUDA kernels by @erwinzhang7 in Add mx.searchsorted with CPU, Metal and CUDA kernels ml-explore/mlx#4035
Stop integer power from hanging on a negative exponent by @ayaangazali in Stop integer power from hanging on a negative exponent ml-explore/mlx#4100
Fix second order gradients for sort, partition, topk and cummax/cummin by @Adityaj0 in Fix second order gradients for sort, partition, topk and cummax/cummin ml-explore/mlx#4117
Fix segfault on negative out of bounds axes in take_along_axis/put_along_axis by @Adityaj0 in Fix segfault on negative out of bounds axes in take_along_axis/put_along_axis ml-explore/mlx#4118
Deduplicate slow CPU convolution dtype dispatch by @PhysicistJohn in Deduplicate slow CPU convolution dtype dispatch ml-explore/mlx#4126
Use common dispatch for compiled constants by @PhysicistJohn in Use common dispatch for compiled constants ml-explore/mlx#4129
Fix set_printoptions example by @PhysicistJohn in Fix set_printoptions example ml-explore/mlx#4140
Clarify export output type diagnostic by @PhysicistJohn in Clarify export output type diagnostic ml-explore/mlx#4141
Fix integer keys collision in tree_unflatten by @Solaris-star in Fix integer keys collision in tree_unflatten ml-explore/mlx#3878
Fix fence tracking for donated dynamic slice offsets by @michaelellis003 in Fix fence tracking for donated dynamic slice offsets ml-explore/mlx#4099
Deduplicate distributed all-reduce construction by @PhysicistJohn in Deduplicate distributed all-reduce construction ml-explore/mlx#4144
Deduplicate distributed socket errors by @PhysicistJohn in Deduplicate distributed socket errors ml-explore/mlx#4145
Remove redundant Metal bitwise dispatch by @PhysicistJohn in Remove redundant Metal bitwise dispatch ml-explore/mlx#4147
Deduplicate custom Metal kernel size validation by @PhysicistJohn in Deduplicate custom Metal kernel size validation ml-explore/mlx#4149
Deduplicate CUBLAS compute type cases by @PhysicistJohn in Deduplicate CUBLAS compute type cases ml-explore/mlx#4151
Refactor CPU Scatter dtype dispatch by @PhysicistJohn in Refactor CPU Scatter dtype dispatch ml-explore/mlx#4155
fix(cpu): name the right routine in LUF and SVD error messages by @devteamaegis in fix(cpu): name the right routine in LUF and SVD error messages ml-explore/mlx#4167
Refactor CPU ArgPartition dtype dispatch by @PhysicistJohn in Refactor CPU ArgPartition dtype dispatch ml-explore/mlx#4156
add complex to scalar to avoid raising warnings by linters by @aaishwarymishra in add complex to scalar to avoid raising warnings by linters ml-explore/mlx#4168
Refactor CPU Copy source dtype dispatch by @PhysicistJohn in Refactor CPU Copy source dtype dispatch ml-explore/mlx#4170
Deduplicate Metal scan encoder setup by @PhysicistJohn in Deduplicate Metal scan encoder setup ml-explore/mlx#4154
[Metal] Skip empty NAX GEMM output groups by @XXXXRT666 in [Metal] Skip empty NAX GEMM output groups ml-explore/mlx#3941
docs: fix stale CLI invocations in LLM inference example by @deBrian07 in docs: fix stale CLI invocations in LLM inference example ml-explore/mlx#4172
Increase ccache size in CI by @zcbenz in Increase ccache size in CI ml-explore/mlx#3999
Fix hadamard_transform on the GPU for n = m with no power-of-2 factor by @axiom-of-choice in Fix hadamard_transform on the GPU for n = m with no power-of-2 factor ml-explore/mlx#4054
Refactor CPU Partition dtype dispatch by @PhysicistJohn in Refactor CPU Partition dtype dispatch ml-explore/mlx#4175
Refactor CPU ScatterAxis and MaskedScatter dtype dispatch by @PhysicistJohn in Refactor CPU ScatterAxis and MaskedScatter dtype dispatch ml-explore/mlx#4176
docs: document softmax's precise argument by @apocryphx in docs: document softmax's precise argument ml-explore/mlx#4178
Send out of range trig arguments to libm by @ayaangazali in Send out of range trig arguments to libm ml-explore/mlx#4157
Add reflect and symmetric padding modes to mx.pad by @katlun-lgtm in Add reflect and symmetric padding modes to mx.pad ml-explore/mlx#3608
docs: add a usage page on growing arrays by @apocryphx in docs: add a usage page on growing arrays ml-explore/mlx#4019
Fix and enable non-transposed NAX qmm by @gordofreemo in Fix and enable non-transposed NAX qmm ml-explore/mlx#4051
Add scatter reduce for JACCL by @angeloskath in Add scatter reduce for JACCL ml-explore/mlx#3901
Fix diag for zero-size input by @devteamaegis in Fix diag for zero-size input ml-explore/mlx#4165
Add an inverse-CDF path to categorical sampling by @deescondrillas in Add an inverse-CDF path to categorical sampling ml-explore/mlx#4177
Fail when a requested RDMA device is not found by @erwinzhang7 in Fail when a requested RDMA device is not found ml-explore/mlx#4180
Bound GGUF tensor data offsets against the file mapping by @robertomeroni in Bound GGUF tensor data offsets against the file mapping ml-explore/mlx#4179
Shift by the max in log_softmax by @ayaangazali in Shift by the max in log_softmax ml-explore/mlx#4169
Shift by the max in cross_entropy by @ayaangazali in Shift by the max in cross_entropy ml-explore/mlx#4188
Fix linalg.norm(x, keepdims=True) not keeping dims by @devteamaegis in Fix linalg.norm(x, keepdims=True) not keeping dims ml-explore/mlx#4166
Refactor inexact CPU binary dtype dispatch by @PhysicistJohn in Refactor inexact CPU binary dtype dispatch ml-explore/mlx#4194
Implement batched matmul for large 1D dot products by @Ved235 in Implement batched matmul for large 1D dot products ml-explore/mlx#3580
Re-disable qmm_n_nax and fix group_size < 64 by @PhilipJohnBasile in Re-disable qmm_n_nax and fix group_size < 64 ml-explore/mlx#4202
fix: shell-quote the working directory in the distributed launch script by @bunlongheng in fix: shell-quote the working directory in the distributed launch script ml-explore/mlx#4189
Report when no usable GID is found by @erwinzhang7 in Report when no usable GID is found ml-explore/mlx#4191
Keep double precision for python floats in float64 operations by @nileshpatil6 in Keep double precision for python floats in float64 operations ml-explore/mlx#4173
Jaccl ring multi wire fix by @angeloskath in Jaccl ring multi wire fix ml-explore/mlx#4193
Report why creating a queue pair failed by @erwinzhang7 in Report why creating a queue pair failed ml-explore/mlx#4209
Add antialias support to nn.Upsample for linear and cubic modes by @lyonsno in Add antialias support to nn.Upsample for linear and cubic modes ml-explore/mlx#3677
Require pytorch 2.12 for metal dlpack tests by @zcbenz in Require pytorch 2.12 for metal dlpack tests ml-explore/mlx#4206
Remove std::optional in C++ public interface by @zcbenz in Remove std::optional<bool> in C++ public interface ml-explore/mlx#4207
Fix SliceUpdate JVPs with one traced input by @PhysicistJohn in Fix SliceUpdate JVPs with one traced input ml-explore/mlx#4200
Raise instead of hanging when reflect/symmetric pad gets an empty axis by @codeAnqiang-ma in Raise instead of hanging when reflect/symmetric pad gets an empty axis ml-explore/mlx#4223
Reuse Stockham twiddles in large batched Bluestein FFTs by @PhysicistJohn in Reuse Stockham twiddles in large batched Bluestein FFTs ml-explore/mlx#4084
Validate decay_steps and step_size in the schedulers by @ayaangazali in Validate decay_steps and step_size in the schedulers ml-explore/mlx#4217
Add runnable distributed examples by @erwinzhang7 in Add runnable distributed examples ml-explore/mlx#4219
Break monolithic MTLResidencySet into smaller sets by @JamesMcCarthy44 in Break monolithic MTLResidencySet into smaller sets ml-explore/mlx#4211
[CUDA][Improvement] RMSNorm backward  by @nastya236 in [CUDA][Improvement] RMSNorm backward  ml-explore/mlx#3881
Fix MultiOptimizer on models containing empty modules by @erenmenges in Fix MultiOptimizer on models containing empty modules ml-explore/mlx#4215
Return tuple in meshgrid by @JasonHonKL in Return tuple in meshgrid ml-explore/mlx#4229
Add endpoint parameter to linspace by @AKnassa in Add endpoint parameter to linspace ml-explore/mlx#4184
Fix vmap of partition/argpartition dropping the kth argument by @Adityaj0 in Fix vmap of partition/argpartition dropping the kth argument ml-explore/mlx#4116
Fix nan_to_num replacing inf with 0 for float16 and bfloat16 by @codeAnqiang-ma in Fix nan_to_num replacing inf with 0 for float16 and bfloat16 ml-explore/mlx#4222
Fix einsum not broadcasting batch dimensions in batched tensordot by @Adityaj0 in Fix einsum not broadcasting batch dimensions in batched tensordot ml-explore/mlx#4125
Dequantize in float32 by @RohanGautam in Dequantize in float32 ml-explore/mlx#4241
Reject complex in erf and erfinv by @ayaangazali in Reject complex in erf and erfinv ml-explore/mlx#4243
Make abs a no-op for unsigned types in the simd helpers by @ayaangazali in Make abs a no-op for unsigned types in the simd helpers ml-explore/mlx#4240
Fix quantize matrix multiplication floor issue by @JasonHonKL in Fix quantize matrix multiplication floor issue ml-explore/mlx#4251
Only use MPI backend for world size > 1 by @erwinzhang7 in Only use MPI backend for world size > 1 ml-explore/mlx#4210
Reject complex in expm1, sigmoid and arctan2 by @ayaangazali in Reject complex in expm1, sigmoid and arctan2 ml-explore/mlx#4257
Decompose small kernel-depth 3D convs into 2D convs by @katlun-lgtm in Decompose small kernel-depth 3D convs into 2D convs ml-explore/mlx#3785
Fix Metal sort of a view with a negative stride by @robertomeroni in Fix Metal sort of a view with a negative stride ml-explore/mlx#4252
Mirror the depth axis in the decomposed 3D conv when flipped by @ayaangazali in Mirror the depth axis in the decomposed 3D conv when flipped ml-explore/mlx#4277
Fix Metal row reductions on negative-stride views by @FU-max-boop in Fix Metal row reductions on negative-stride views ml-explore/mlx#4267
[CUDA] Fix custom kernel cache collision for same name, different source by @HaoXuAI in [CUDA] Fix custom kernel cache collision for same name, different source ml-explore/mlx#4273
Fix ops rejecting integers larger than INT32_MAX by @FeliGame in Fix ops rejecting integers larger than INT32_MAX ml-explore/mlx#4255
Fix var/std for complex numbers by @ayaangazali in Fix var/std for complex numbers ml-explore/mlx#4260
Fix int32 overflow in conv padded input and pad shapes by @eyupcanakman in Fix int32 overflow in conv padded input and pad shapes ml-explore/mlx#4258
Reject complex in remainder by @ayaangazali in Reject complex in remainder ml-explore/mlx#4270
Compare the macOS SDK version as a version when gating JACCL by @erwinzhang7 in Compare the macOS SDK version as a version when gating JACCL ml-explore/mlx#4286
Clamp ring socket transfers so a payload of 2 GiB or more can be sent by @erwinzhang7 in Clamp ring socket transfers so a payload of 2 GiB or more can be sent ml-explore/mlx#4281
Replace hand-rolled axis normalization with normalize_axis_index in split/unstack/partition/topk by @Adityaj0 in Replace hand-rolled axis normalization with normalize_axis_index in split/unstack/partition/topk ml-explore/mlx#4288
Remove grouped output in CI by @zcbenz in Remove grouped output in CI ml-explore/mlx#4195
[CUDA] Fix finding cuda 13 headers in JIT compilation by @zcbenz in [CUDA] Fix finding cuda 13 headers in JIT compilation ml-explore/mlx#3995
Refactor wheel building script by @zcbenz in Refactor wheel building script ml-explore/mlx#3818
Make mx.compile cache erasing thread safe by @zcbenz in Make mx.compile cache erasing thread safe ml-explore/mlx#4248
Add builds for free-threaded python by @zcbenz in Add builds for free-threaded python ml-explore/mlx#3812
Fix int32 overflow in concatenate/repeat/kron by @ayaangazali in Fix int32 overflow in concatenate/repeat/kron ml-explore/mlx#4303
Widen list elements that do not fit in int32 to int64 by @ayaangazali in Widen list elements that do not fit in int32 to int64 ml-explore/mlx#4305
Propagate CPU errors to events by @zcbenz in  Propagate CPU errors to events ml-explore/mlx#3742
Fix mx.arange dtype inference overflow regression by @shoutoutuoadi325 in Fix mx.arange dtype inference overflow regression ml-explore/mlx#4324
Add workflow to update pull request limit bypass list by @zcbenz in Add workflow to update pull request limit bypass list ml-explore/mlx#4320
Support head dimension 72 in Metal full attention by @dhiltgen in Support head dimension 72 in Metal full attention ml-explore/mlx#4330

New Contributors

@pierre427 made their first contribution in Warn at configure time when NAX kernels are disabled ml-explore/mlx#3824
@apocryphx made their first contribution in Fix JIT preamble header filter matching project paths containing "Xcode" ml-explore/mlx#3873
@ayaangazali made their first contribution in Document default value of p in dropout layer docstrings ml-explore/mlx#3870
@HaoXuAI made their first contribution in Zero-copy CPU import: mx.array(host_buffer, copy=False) on unified memory ml-explore/mlx#3872
@neilmehta24 made their first contribution in Reuse Metal WAR tracking hash tables ml-explore/mlx#3882
@wyanzhao made their first contribution in Use unroll_count(4) for the NAX attention Q@K.T loop ml-explore/mlx#3843
@metascroy made their first contribution in Fix incorrect nvfp4 quantized_matmul through the split-K path ml-explore/mlx#3854
@aleroot made their first contribution in [Metal] Avoid regex in custom kernel name generation ml-explore/mlx#3869
@AdamDLuz made their first contribution in [CUDA] Fix grid overflow in gemm conv unfold kernels for >= 65,536 output positions ml-explore/mlx#3893
@vineethsaivs made their first contribution in Fix Glorot/He uniform init docstrings to label the uniform bound, not sigma ml-explore/mlx#3831
@winklemad made their first contribution in Fix log_cosh_loss docstring to document the element-wise loss ml-explore/mlx#3846
@stoyoda0012-cyber made their first contribution in docs: document the reduced-precision float32 default and MLX_ENABLE_TF32 ml-explore/mlx#3894
@PhysicistJohn made their first contribution in Export C++20 requirement to CMake consumers ml-explore/mlx#3971
@latent-9 made their first contribution in docs: remove references to removed --no-verify-script launch flag ml-explore/mlx#3959
@aaishwarymishra made their first contribution in added eye(0) support and tests ml-explore/mlx#3952
@reckylurker made their first contribution in python: fix bfloat16 buffer format itemsize mismatch ml-explore/mlx#3975
@sashko-zakharchuk made their first contribution in Fix mx.remainder floored-mod for float16/bfloat16 on CPU ml-explore/mlx#3976
@WindChimeRan made their first contribution in Fix state corruption when a primitive throws during eval ml-explore/mlx#3675
@jonathan308 made their first contribution in Fix mlx.launch --python: flag is parsed but never forwarded to the launch script ml-explore/mlx#4002
@jasp-nerd made their first contribution in Fix CUDA batched GEMV grid overflow ml-explore/mlx#3929
@Gusanidas made their first contribution in Fix segfault in expand_dims for out of bounds negative axes ml-explore/mlx#4021
@JasonHonKL made their first contribution in Fix mx.longsumexp output shape issue. ml-explore/mlx#4030
@yingjiacai made their first contribution in Fix installed static MLX package on Windows ml-explore/mlx#3848
@jlee5814 made their first contribution in Derive the qmv fast path K alignment from bits ml-explore/mlx#3965
@kapellirohith made their first contribution in Fix eigh UPLO and zero-size eigh/svd on the CPU ml-explore/mlx#3834
@prady0t made their first contribution in Fix empty matrix case in cholesky ml-explore/mlx#4033
@magnuslundstedt made their first contribution in Request MSL 4.1 from the runtime compiler on macOS 27 ml-explore/mlx#4052
@kitty-eu-org made their first contribution in Fix concurrent Metal kernel cache lookup ml-explore/mlx#4043
@dudududukim made their first contribution in Normalize biases before encoding in gather_qmm_rhs ml-explore/mlx#4056
@erwinzhang7 made their first contribution in ring: fail on peer disconnect instead of hanging forever ml-explore/mlx#4060
@YassineMA03 made their first contribution in Fix SinusoidalPositionalEncoding silently overriding explicit scale=0.0 ml-explore/mlx#4098
@dwijenpatel made their first contribution in Pick BM from rows per expert in gather_qmm_rhs_nax ml-explore/mlx#4023
@Adityaj0 made their first contribution in Fix second order gradients for sort, partition, topk and cummax/cummin ml-explore/mlx#4117
@Solaris-star made their first contribution in Fix integer keys collision in tree_unflatten ml-explore/mlx#3878
@michaelellis003 made their first contribution in Fix fence tracking for donated dynamic slice offsets ml-explore/mlx#4099
@devteamaegis made their first contribution in fix(cpu): name the right routine in LUF and SVD error messages ml-explore/mlx#4167
@deBrian07 made their first contribution in docs: fix stale CLI invocations in LLM inference example ml-explore/mlx#4172
@axiom-of-choice made their first contribution in Fix hadamard_transform on the GPU for n = m with no power-of-2 factor ml-explore/mlx#4054
@gordofreemo made their first contribution in Fix and enable non-transposed NAX qmm ml-explore/mlx#4051
@deescondrillas made their first contribution in Add an inverse-CDF path to categorical sampling ml-explore/mlx#4177
@robertomeroni made their first contribution in Bound GGUF tensor data offsets against the file mapping ml-explore/mlx#4179
@Ved235 made their first contribution in Implement batched matmul for large 1D dot products ml-explore/mlx#3580
@PhilipJohnBasile made their first contribution in Re-disable qmm_n_nax and fix group_size < 64 ml-explore/mlx#4202
@bunlongheng made their first contribution in fix: shell-quote the working directory in the distributed launch script ml-explore/mlx#4189
@nileshpatil6 made their first contribution in Keep double precision for python floats in float64 operations ml-explore/mlx#4173
@codeAnqiang-ma made their first contribution in Raise instead of hanging when reflect/symmetric pad gets an empty axis ml-explore/mlx#4223
@JamesMcCarthy44 made their first contribution in Break monolithic MTLResidencySet into smaller sets ml-explore/mlx#4211
@erenmenges made their first contribution in Fix MultiOptimizer on models containing empty modules ml-explore/mlx#4215
@AKnassa made their first contribution in Add endpoint parameter to linspace ml-explore/mlx#4184
@RohanGautam made their first contribution in Dequantize in float32 ml-explore/mlx#4241
@FU-max-boop made their first contribution in Fix Metal row reductions on negative-stride views ml-explore/mlx#4267
@FeliGame made their first contribution in Fix ops rejecting integers larger than INT32_MAX ml-explore/mlx#4255
@shoutoutuoadi325 made their first contribution in Fix mx.arange dtype inference overflow regression ml-explore/mlx#4324

Full Changelog: ml-explore/mlx@v0.32.0...v0.32.1

View the full release notes at https://github.com/ml-explore/mlx/releases/tag/v0.32.1.